#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_X, typename TYPE_Y>
class MishKernel {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> xQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;

    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> Cal1, Cal2, Tmp1, Tmp2;
    // TBuf<QuePosition::VECCALC> xBuf, yBuf;
    // TBuf<QuePosition::VECCALC> Tmp_F, Row_F, Col_F, Lidx_F, Gidx_F;

    uint32_t ALIGN_NUM;
    uint32_t totalLength;
    uint32_t tileLength;
    uint32_t tileNum;
    uint32_t lasttileLength;

public:
    __aicore__ inline MishKernel() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t ALIGN_NUM, uint32_t totalLength,
                                uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength)
    {
        this->ALIGN_NUM = ALIGN_NUM;
        this->totalLength = totalLength;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;

        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y, totalLength);


        pipe.InitBuffer(xQueue, BUFFER_NUM, tileLength * sizeof(TYPE_X)); //自动向上32Byte对齐
        pipe.InitBuffer(yQueue, BUFFER_NUM, tileLength * sizeof(TYPE_Y));

        pipe.InitBuffer(Cal1, tileLength * sizeof(TYPE_Y));
        pipe.InitBuffer(Cal2, tileLength * sizeof(TYPE_Y));
        // pipe.InitBuffer(Cal3, tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(Tmp1, tileLength * sizeof(TYPE_Y));
        pipe.InitBuffer(Tmp2, tileLength * sizeof(TYPE_Y));
        // pipe.InitBuffer(Tmp3, tileLength * sizeof(TYPE_X));
        // pipe.InitBuffer(xBuf, tileLength * sizeof(float));
        // pipe.InitBuffer(yBuf, tileLength * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        uint32_t  loopCount = tileNum;
        
        for (int32_t i = 0; i < loopCount - 1; i++) 
        {                
            CopyIn(i, tileLength);
            Compute(i, tileLength);
            CopyOut(i, tileLength);
        }

        CopyIn(tileNum - 1, lasttileLength);
        Compute(tileNum - 1, lasttileLength);
        CopyOut(tileNum - 1, lasttileLength);
            
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = xQueue.AllocTensor<TYPE_X>();
        uint32_t startindex = progress * tileLength;
        DataCopy(xLocal, xGm[startindex], length);
        xQueue.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = xQueue.DeQue<TYPE_X>();
        LocalTensor<TYPE_Y> yLocal = yQueue.AllocTensor<TYPE_Y>();

        LocalTensor<TYPE_X> cal1 = Cal1.Get<TYPE_X>();
        LocalTensor<TYPE_X> cal2 = Cal2.Get<TYPE_X>();
        
        LocalTensor<TYPE_X> tmp1 = Tmp1.Get<TYPE_X>();
        LocalTensor<TYPE_X> tmp2 = Tmp2.Get<TYPE_X>();

        // LocalTensor<float> xLocal_f = xBuf.Get<float>();
        // LocalTensor<float> yLocal_f = yBuf.Get<float>();

        // if constexpr(std::is_same_v<TYPE_X, half>)
        // {
        //     Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
        // }
        // else
        // {
        //     xLocal_f = xLocal;
        // }


        Exp(cal1, xLocal, length);
        Adds(cal1, cal1, static_cast<TYPE_X>(1.0f), length);
        Ln(cal1, cal1, length); //

        Muls(cal1, cal1, static_cast<TYPE_X>(2.0f), length);  // 2ln(1+e^x)

        Exp(cal2, cal1, length);    //e^2ln(1+e^x)

        // Adds(tmp1, cal2, static_cast<TYPE_X>(-1.0), length);  //e^2ln(1+e^x) - 1
        Adds(tmp2, cal2, static_cast<TYPE_X>(1.0f), length);  //e^2ln(1+e^x) + 1

        Muls(tmp1, xLocal, static_cast<TYPE_X>(-2.0f), length); // -2x

        Div(tmp1, tmp1, tmp2, length);  // -2x/(e^2ln(1+e^x) + 1)
        Add(yLocal, xLocal, tmp1, length);  // x - 2x/(e^2ln(1+e^x) + 1)

        // if constexpr(std::is_same_v<TYPE_X, half>)
        // {
        //     Cast(yLocal, yLocal_f, RoundMode::CAST_NONE, length);
        // }
        // else
        // {
        //     yLocal = yLocal_f;
        // }

        xQueue.FreeTensor(xLocal);
        yQueue.EnQue<DTYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_Y> yLocal = yQueue.DeQue<DTYPE_Y>();
        uint32_t startindex = progress * tileLength;
        DataCopy(yGm[startindex], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }
};

extern "C" __global__ __aicore__ void mish(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    MishKernel<DTYPE_X, DTYPE_Y> op;
    op.Init(x, y, tiling_data.ALIGN_NUM, tiling_data.totalLength,
            tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength);
    op.Process();
}